package com.sliding.aspectj;

import com.common.utils.IpUtils;
import com.common.utils.ServletUtils;
import com.sliding.annotation.RateLimitByLua;
import com.sliding.enums.ModeType;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.List;

@Aspect //告诉框架我是一个切面类
@Component// 随着框架的启动而启动
@Slf4j
@AllArgsConstructor
public class RateLimitByLuaAspect {

    private final RedisTemplate<Object, Object> redisTemplate;

    private RedisScript<Long> limitScript;

    @Autowired
    public void setLimitScript(RedisScript<Long> limitScript) {
        this.limitScript = limitScript;
    }

    @Around("@annotation(rateLimitByLua)")
    public Object around(ProceedingJoinPoint point, RateLimitByLua rateLimitByLua) throws Throwable {
        // 组合限流 key
        String limitKey = generateKey(point, rateLimitByLua);
        // 执行限流逻辑
        boolean limited = isLimited(limitKey, rateLimitByLua.window(), rateLimitByLua.threshold());
        if (limited) {
            return rateLimitByLua.message();
        }
        // 执行目标方法
        return point.proceed();
    }


    private String generateKey(ProceedingJoinPoint point, RateLimitByLua rateLimitByLua) {
        StringBuilder sb = new StringBuilder(rateLimitByLua.prefix());
        // 获取请求的类名
        String className = point.getTarget().getClass().getName();
        // 获取请求的方法名
        String methodName = point.getSignature().getName();
        if (rateLimitByLua.mode() == ModeType.IP) {
            sb.append(IpUtils.getIpAddress(ServletUtils.getRequest())).append("-");
        }
        sb.append(className).append("-").append(methodName);
        return sb.toString();
    }


    private boolean isLimited(String key, int window, int threshold) {
        List<Object> keys = Collections.singletonList(key);
        try {
            //剩余可用请求数，-1表示被限流
            Long remainingCount = redisTemplate.execute(limitScript, keys, String.valueOf(window), String.valueOf(threshold), String.valueOf(System.currentTimeMillis()));
            log.info("限制请求'{}',剩余请求'{}',缓存key'{}'", threshold, remainingCount.intValue(), key);
            return (remainingCount == null ? -1 : remainingCount) < 0;
        } catch (Exception e) {
            log.error("Redis rate limiter error", e);
            // 发生异常时放行请求
            return true;
        }
    }
}
